package main

import (
	"bufio"
	"context"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"

	cdp "github.com/chromedp/chromedp"
)

func main() {
	// var geturl []string
	geturl, err := getaddress()
	if err != nil {
		println("getaddress is fail")
	}
	dirName := strconv.FormatInt(time.Now().Unix(), 10)
	println("诚享智能试题：")
	os.Mkdir(dirName, os.ModePerm)
	// println(" geturl size %d %d", len(geturl), cap(geturl))
	for key, val := range geturl {
		// fmt.Fprint("id = %d ,url = %s", key, val)
		fmt.Printf("id = %d ,url = %s \n", key, val)
		getwebs(key, val, dirName)
	}

}

func getaddress() ([]string, error) {
	rf, err := os.Open("address.txt")
	if err != nil {
		fmt.Println("请检查当前目录下address.txt文件是否存在，并且内容及格式是否正确！")
		return nil, errors.New("file is err!")
	}
	buf := bufio.NewReader(rf)
	var urls []string
	for {
		line, err := buf.ReadString('\n')
		line = strings.TrimSpace(line)
		if err != nil {
			if err == io.EOF {
				// fmt.Printf(" s = %s", line) //若最后一行没有加换行符，则程序读到最后一行读到EOF，则认为此行为结束行。
				if line != "" {
					urls = append(urls, line)
				}
				break
			}
			fmt.Println("读取文件出错！")
			break
		}
		if line == "" {
			continue
		}
		urls = append(urls, line)
	}
	// for key, val := range urls {
	// 	// fmt.Fprint("id = %d ,url = %s", key, val)
	// 	fmt.Printf("id = %d ,url = %s \n", key, val)
	// }
	// fmt.Println("\ntest")
	return urls, nil
}

func getwebs(urlid int, url, dirName string) {
	response, err := http.Get(url)
	if err != nil {
		//...
		fmt.Println("该地址处理失败:", err)
		return
	}
	defer response.Body.Close() //在回复后必须关闭回复的主体
	// 创建新的cdp上下文
	ctx, cancel := cdp.NewContext(context.Background())
	defer cancel()

	// 此处以首页为例
	// urlstr := `https://codeforces.com/problemset/problem/1658/A`
	var buf []byte
	var filename string

	// 需要截图的元素，支持CSS selector以及XPath query
	//   selector := `#main`
	selector := `#pageContent`
	// selector := `#body`
	if err := cdp.Run(ctx, elementScreen(url, selector, &buf, &filename)); err != nil {
		log.Fatal(err)
	}
	filename = dirName + "\\" + strconv.Itoa(urlid) + filename[3:] + ".png"
	println("文件存放目录：", filename)
	// 写入文件
	if err := ioutil.WriteFile(filename, buf, 0644); err != nil {
		log.Fatal(err)
	}

}

// 截图方法
func elementScreen(urlstr, sel string, res *[]byte, filename *string) cdp.Tasks {
	return cdp.Tasks{
		// 打开url指向的页面
		cdp.Navigate(urlstr),

		// 也可以等待一定的时间
		cdp.Sleep(time.Duration(3) * time.Second),
		// cdp.Text(`.title`, &filename, cdp.NodeVisible),
		cdp.Text(`.problem-statement .header .title`, filename, cdp.NodeVisible),
		// 等待待截图的元素渲染完成
		cdp.WaitVisible(sel, cdp.ByID),

		// 执行截图
		cdp.Screenshot(sel, res, cdp.NodeVisible, cdp.ByID),
	}
}
